package jamirr.apt;

import jamirr.Mirror;
import jamirr.MirrorAnnotation;

import javax.annotation.processing.AbstractProcessor;
import javax.annotation.processing.RoundEnvironment;
import javax.lang.model.SourceVersion;
import javax.lang.model.element.AnnotationMirror;
import javax.lang.model.element.Element;
import javax.lang.model.element.ElementKind;
import javax.lang.model.element.TypeElement;
import javax.lang.model.type.DeclaredType;
import javax.tools.Diagnostic;
import java.io.PrintWriter;
import java.io.StringWriter;
import java.util.Collections;
import java.util.List;
import java.util.Set;

public abstract class AbstractElementProcessor extends AbstractProcessor {

    @Override
    public Set<String> getSupportedAnnotationTypes() {
        return Collections.singleton("*");
    }

    @Override
    public SourceVersion getSupportedSourceVersion() {
        return SourceVersion.latestSupported();
    }

    @Override
    public final boolean process(Set<? extends TypeElement> annotationTypes, RoundEnvironment roundEnv) {
        if (roundEnv.processingOver()) {
            return false;
        }

        var ctx = new ProcessorContext(
                processingEnv.getElementUtils(),
                processingEnv.getTypeUtils(),
                processingEnv.getFiler()
        );

        try {
            Set<? extends Element> elements = roundEnv.getElementsAnnotatedWith(MirrorAnnotation.class);
            for (Element element : elements) {
                if (element.getKind() == ElementKind.ANNOTATION_TYPE) {
                    processAnnotation((DeclaredType) element.asType(), ctx);
                    ctx.addMirrorAnnotation(element.toString());
                }
            }

            for (TypeElement annotationType : annotationTypes) {
                if (isMirrorAnnotated((DeclaredType) annotationType.asType(), ctx)) {
                    Set<? extends Element> enclosedElements = roundEnv.getElementsAnnotatedWith(annotationType);
                    for (Element enclosedElement : enclosedElements) {
                        if (enclosedElement.getKind() == ElementKind.CLASS) {
                            processElement((DeclaredType) enclosedElement.asType(), ctx);
                        }
                    }
                }
            }
        } catch (RuntimeException ex) {
            var sw = new StringWriter();
            var pw = new PrintWriter(sw);
            ex.printStackTrace(pw);
            processingEnv.getMessager().printMessage(Diagnostic.Kind.ERROR, sw.toString());
        }

        return true;
    }

    protected abstract void processAnnotation(DeclaredType declaredType, ProcessorContext ctx);

    protected abstract void processElement(DeclaredType declaredType, ProcessorContext ctx);

    private boolean isMirrorAnnotated(DeclaredType declaredType, ProcessorContext ctx) {
        if (ctx.isInternalAnnotation(declaredType)) {
            return false;
        }
        if (ctx.isSame(declaredType, Mirror.class.getCanonicalName())) {
            return true;
        }

        if (!ctx.isTypeElementAnnotated(declaredType)) {
            return false;
        }

        List<? extends AnnotationMirror> annotationMirrors = declaredType.asElement().getAnnotationMirrors();
        return annotationMirrors.stream()
                .anyMatch(a -> isMirrorAnnotated(a.getAnnotationType(), ctx));
    }
}
